{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4c599bc9",
   "metadata": {},
   "source": [
    "# Tutorial 06, case 5: Stokes problem with distributed control\n",
    "\n",
    "In this tutorial we solve the optimal control problem\n",
    "\n",
    "$$\\min J(y, u) = \\frac{1}{2} \\int_{\\Omega} |v - v_d|^2 dx + \\frac{\\alpha}{2} \\int_{\\Omega} |u|^2 dx$$\n",
    "s.t.\n",
    "$$\\begin{cases}\n",
    "- \\Delta v + \\nabla p = f + u   & \\text{in } \\Omega\\\\\n",
    "         \\text{div} v = 0       & \\text{in } \\Omega\\\\\n",
    "                    v = 0       & \\text{on } \\partial\\Omega\n",
    "\\end{cases}$$\n",
    "\n",
    "where\n",
    "$$\\begin{align*}\n",
    "& \\Omega                      & \\text{unit square}\\\\\n",
    "& u \\in [L^2(\\Omega)]^2       & \\text{control variable}\\\\\n",
    "& v \\in [H^1_0(\\Omega)]^2     & \\text{state velocity variable}\\\\\n",
    "& p \\in L^2(\\Omega)           & \\text{state pressure variable}\\\\\n",
    "& \\alpha > 0                  & \\text{penalization parameter}\\\\\n",
    "& v_d                         & \\text{desired state}\\\\\n",
    "& f                           & \\text{forcing term}\n",
    "\\end{align*}$$\n",
    "using an adjoint formulation solved by a one shot approach"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c531f0b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dolfinx.fem\n",
    "import dolfinx.fem.petsc\n",
    "import dolfinx.io\n",
    "import dolfinx.mesh\n",
    "import mpi4py.MPI\n",
    "import numpy as np\n",
    "import numpy.typing\n",
    "import petsc4py.PETSc\n",
    "import sympy\n",
    "import ufl\n",
    "import viskex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "779a06f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import multiphenicsx.fem\n",
    "import multiphenicsx.fem.petsc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f57e807",
   "metadata": {},
   "source": [
    "### Mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0af49590",
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh = dolfinx.mesh.create_unit_square(mpi4py.MPI.COMM_WORLD, 32, 32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "167dc814-3ca7-4aee-aa0f-fccf2f477e32",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create connectivities required by the rest of the code\n",
    "mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d6f60c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def bottom(x: np.typing.NDArray[np.float64]) -> np.typing.NDArray[np.bool_]:\n",
    "    \"\"\"Condition that defines the bottom boundary.\"\"\"\n",
    "    return abs(x[1] - 0.) < np.finfo(float).eps  # type: ignore[no-any-return]\n",
    "\n",
    "\n",
    "def left(x: np.typing.NDArray[np.float64]) -> np.typing.NDArray[np.bool_]:\n",
    "    \"\"\"Condition that defines the left boundary.\"\"\"\n",
    "    return abs(x[0] - 0.) < np.finfo(float).eps  # type: ignore[no-any-return]\n",
    "\n",
    "\n",
    "def top(x: np.typing.NDArray[np.float64]) -> np.typing.NDArray[np.bool_]:\n",
    "    \"\"\"Condition that defines the top boundary.\"\"\"\n",
    "    return abs(x[1] - 1.) < np.finfo(float).eps  # type: ignore[no-any-return]\n",
    "\n",
    "\n",
    "def right(x: np.typing.NDArray[np.float64]) -> np.typing.NDArray[np.bool_]:\n",
    "    \"\"\"Condition that defines the right boundary.\"\"\"\n",
    "    return abs(x[0] - 1.) < np.finfo(float).eps  # type: ignore[no-any-return]\n",
    "\n",
    "\n",
    "boundaries_entities = dict()\n",
    "boundaries_values = dict()\n",
    "for (boundary, boundary_id) in zip((bottom, left, top, right), (1, 2, 3, 4)):\n",
    "    boundaries_entities[boundary_id] = dolfinx.mesh.locate_entities_boundary(\n",
    "        mesh, mesh.topology.dim - 1, boundary)\n",
    "    boundaries_values[boundary_id] = np.full(\n",
    "        boundaries_entities[boundary_id].shape, boundary_id, dtype=np.int32)\n",
    "boundaries_entities_unsorted = np.hstack(list(boundaries_entities.values()))\n",
    "boundaries_values_unsorted = np.hstack(list(boundaries_values.values()))\n",
    "boundaries_entities_argsort = np.argsort(boundaries_entities_unsorted)\n",
    "boundaries_entities_sorted = boundaries_entities_unsorted[boundaries_entities_argsort]\n",
    "boundaries_values_sorted = boundaries_values_unsorted[boundaries_entities_argsort]\n",
    "boundaries = dolfinx.mesh.meshtags(\n",
    "    mesh, mesh.topology.dim - 1,\n",
    "    boundaries_entities_sorted, boundaries_values_sorted)\n",
    "boundaries.name = \"boundaries\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bab45f7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "boundaries_1234 = boundaries.indices[np.isin(boundaries.values, (1, 2, 3, 4))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1b1d2b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_mesh(mesh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3b74e5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_mesh_tags(mesh, boundaries, \"boundaries\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ac5a0cf",
   "metadata": {},
   "source": [
    "### Function spaces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d9a6e35",
   "metadata": {},
   "outputs": [],
   "source": [
    "Y_velocity = dolfinx.fem.functionspace(mesh, (\"Lagrange\", 2, (mesh.geometry.dim, )))\n",
    "Y_pressure = dolfinx.fem.functionspace(mesh, (\"Lagrange\", 1))\n",
    "U = dolfinx.fem.functionspace(mesh, (\"Lagrange\", 2, (mesh.geometry.dim, )))\n",
    "Q_velocity = Y_velocity.clone()\n",
    "Q_pressure = Y_pressure.clone()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0f2b2d6",
   "metadata": {},
   "source": [
    "### Trial and test functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f68b72c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "(v, p) = (ufl.TrialFunction(Y_velocity), ufl.TrialFunction(Y_pressure))\n",
    "(w, q) = (ufl.TestFunction(Y_velocity), ufl.TestFunction(Y_pressure))\n",
    "u = ufl.TrialFunction(U)\n",
    "r = ufl.TestFunction(U)\n",
    "(z, b) = (ufl.TrialFunction(Q_velocity), ufl.TrialFunction(Q_pressure))\n",
    "(s, d) = (ufl.TestFunction(Q_velocity), ufl.TestFunction(Q_pressure))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07a6aac1",
   "metadata": {},
   "source": [
    " ### Problem data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c622a593",
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha = 1.e-5\n",
    "epsilon = 1.e-5\n",
    "x, y = sympy.symbols(\"x[0], x[1]\")\n",
    "psi_d = 10 * (1 - sympy.cos(0.8 * np.pi * x)) * (1 - sympy.cos(0.8 * np.pi * y)) * (1 - x)**2 * (1 - y)**2\n",
    "v_d_x = sympy.lambdify([x, y], psi_d.diff(y, 1))\n",
    "v_d_y = sympy.lambdify([x, y], - psi_d.diff(x, 1))\n",
    "v_d = dolfinx.fem.Function(Y_velocity)\n",
    "v_d.interpolate(lambda x: np.stack((v_d_x(x[0], x[1]), v_d_y(x[0], x[1])), axis=0))\n",
    "ff = dolfinx.fem.Constant(mesh, tuple(petsc4py.PETSc.ScalarType(0) for _ in range(2)))\n",
    "bc0 = np.zeros((2, ), dtype=petsc4py.PETSc.ScalarType)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4acc64b4",
   "metadata": {},
   "source": [
    "### Optimality conditions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "192f72ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "a = [[ufl.inner(v, w) * ufl.dx, None, None, ufl.inner(ufl.grad(z), ufl.grad(w)) * ufl.dx,\n",
    "      - ufl.inner(b, ufl.div(w)) * ufl.dx],\n",
    "     [None, None, None, - ufl.inner(ufl.div(z), q) * ufl.dx, epsilon * ufl.inner(b, q) * ufl.dx],\n",
    "     [None, None, alpha * ufl.inner(u, r) * ufl.dx, - ufl.inner(z, r) * ufl.dx, None],\n",
    "     [ufl.inner(ufl.grad(v), ufl.grad(s)) * ufl.dx, - ufl.inner(p, ufl.div(s)) * ufl.dx,\n",
    "      - ufl.inner(u, s) * ufl.dx, None, None],\n",
    "     [- ufl.inner(ufl.div(v), d) * ufl.dx, epsilon * ufl.inner(p, d) * ufl.dx, None, None, None]]\n",
    "f = [ufl.inner(v_d, w) * ufl.dx,\n",
    "     None,\n",
    "     None,\n",
    "     ufl.inner(ff, s) * ufl.dx,\n",
    "     None]\n",
    "a[3][3] = dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(0)) * ufl.inner(z, s) * ufl.dx\n",
    "f[1] = ufl.inner(dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(0)), q) * ufl.dx\n",
    "f[2] = ufl.inner(dolfinx.fem.Constant(mesh, tuple(petsc4py.PETSc.ScalarType(0) for _ in range(2))), r) * ufl.dx\n",
    "f[4] = ufl.inner(dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(0)), d) * ufl.dx\n",
    "a_cpp = dolfinx.fem.form(a)\n",
    "f_cpp = dolfinx.fem.form(f)\n",
    "bdofs_Y_velocity_1234 = dolfinx.fem.locate_dofs_topological(\n",
    "    Y_velocity, mesh.topology.dim - 1, boundaries_1234)\n",
    "bdofs_Q_velocity_1234 = dolfinx.fem.locate_dofs_topological(\n",
    "    Q_velocity, mesh.topology.dim - 1, boundaries_1234)\n",
    "bc = [dolfinx.fem.dirichletbc(bc0, bdofs_Y_velocity_1234, Y_velocity),\n",
    "      dolfinx.fem.dirichletbc(bc0, bdofs_Q_velocity_1234, Q_velocity)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ace7eea",
   "metadata": {},
   "source": [
    "### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce90efd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "(v, p) = (dolfinx.fem.Function(Y_velocity), dolfinx.fem.Function(Y_pressure))\n",
    "u = dolfinx.fem.Function(U)\n",
    "(z, b) = (dolfinx.fem.Function(Q_velocity), dolfinx.fem.Function(Q_pressure))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "008a25ef",
   "metadata": {},
   "source": [
    "### Cost functional"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ea13f4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "J = 0.5 * ufl.inner(v - v_d, v - v_d) * ufl.dx + 0.5 * alpha * ufl.inner(u, u) * ufl.dx\n",
    "J_cpp = dolfinx.fem.form(J)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22e908e2",
   "metadata": {},
   "source": [
    "### Uncontrolled functional value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08b4ee22",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract state forms from the optimality conditions\n",
    "a_state = [[ufl.replace(a[i][j], {s: w, d: q}) if a[i][j] is not None else None\n",
    "            for j in (0, 1)] for i in (3, 4)]\n",
    "f_state = [ufl.replace(f[i], {s: w, d: q}) for i in (3, 4)]\n",
    "a_state_cpp = dolfinx.fem.form(a_state)\n",
    "f_state_cpp = dolfinx.fem.form(f_state)\n",
    "bc_state = [bc[0]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "116f4b75",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assemble the block linear system for the state\n",
    "A_state = dolfinx.fem.petsc.assemble_matrix_block(a_state_cpp, bcs=bc_state)\n",
    "A_state.assemble()\n",
    "F_state = dolfinx.fem.petsc.assemble_vector_block(f_state_cpp, a_state_cpp, bcs=bc_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f61a3a1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve\n",
    "vp = dolfinx.fem.petsc.create_vector_block([f_cpp[j] for j in (0, 1)])\n",
    "ksp = petsc4py.PETSc.KSP()\n",
    "ksp.create(mesh.comm)\n",
    "ksp.setOperators(A_state)\n",
    "ksp.setType(\"preonly\")\n",
    "ksp.getPC().setType(\"lu\")\n",
    "ksp.getPC().setFactorSolverType(\"mumps\")\n",
    "ksp.setFromOptions()\n",
    "ksp.solve(F_state, vp)\n",
    "vp.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)\n",
    "ksp.destroy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5ea9d74",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split the block solution in components\n",
    "with multiphenicsx.fem.petsc.BlockVecSubVectorWrapper(vp, [c.function_space.dofmap for c in (v, p)]) as vp_wrapper:\n",
    "    for vp_wrapper_local, component in zip(vp_wrapper, (v, p)):\n",
    "        with component.x.petsc_vec.localForm() as component_local:\n",
    "            component_local[:] = vp_wrapper_local"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4f21ebc",
   "metadata": {},
   "outputs": [],
   "source": [
    "J_uncontrolled = mesh.comm.allreduce(dolfinx.fem.assemble_scalar(J_cpp), op=mpi4py.MPI.SUM)\n",
    "print(\"Uncontrolled J =\", J_uncontrolled)\n",
    "assert np.isclose(J_uncontrolled, 0.1784536)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbe65be0",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_vector_field(v, \"uncontrolled state velocity\", glyph_factor=1e-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc7baefa",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(p, \"uncontrolled state pressure\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4bbea05",
   "metadata": {},
   "source": [
    "### Optimal control"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08f81c12",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assemble the block linear system for the optimality conditions\n",
    "A = dolfinx.fem.petsc.assemble_matrix_block(a_cpp, bcs=bc)\n",
    "A.assemble()\n",
    "F = dolfinx.fem.petsc.assemble_vector_block(f_cpp, a_cpp, bcs=bc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93e25eb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve\n",
    "vpuzb = dolfinx.fem.petsc.create_vector_block(f_cpp)\n",
    "ksp = petsc4py.PETSc.KSP()\n",
    "ksp.create(mesh.comm)\n",
    "ksp.setOperators(A)\n",
    "ksp.setType(\"preonly\")\n",
    "ksp.getPC().setType(\"lu\")\n",
    "ksp.getPC().setFactorSolverType(\"mumps\")\n",
    "ksp.setFromOptions()\n",
    "ksp.solve(F, vpuzb)\n",
    "vpuzb.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)\n",
    "ksp.destroy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5966c006",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split the block solution in components\n",
    "with multiphenicsx.fem.petsc.BlockVecSubVectorWrapper(\n",
    "        vpuzb, [c.function_space.dofmap for c in (v, p, u, z, b)]) as vpuzb_wrapper:\n",
    "    for vpuzb_wrapper_local, component in zip(vpuzb_wrapper, (v, p, u, z, b)):\n",
    "        with component.x.petsc_vec.localForm() as component_local:\n",
    "            component_local[:] = vpuzb_wrapper_local"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf67db10",
   "metadata": {},
   "outputs": [],
   "source": [
    "J_controlled = mesh.comm.allreduce(dolfinx.fem.assemble_scalar(J_cpp), op=mpi4py.MPI.SUM)\n",
    "print(\"Optimal J =\", J_controlled)\n",
    "assert np.isclose(J_controlled, 0.0052941)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fbf3d8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_vector_field(v, \"state velocity\", glyph_factor=1e-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1de6a28",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(p, \"state pressure\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd17e203",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_vector_field(u, \"control\", glyph_factor=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e4b7f92",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_vector_field(z, \"adjoint velocity\", glyph_factor=1e2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe39b51a",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(b, \"adjoint pressure\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython"
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
